"""An example of embedding a RichJupyterWidget in a PyQT Application.
This uses a normal kernel launched as a subprocess. It shows how to shutdown
the kernel cleanly when the application quits.
To run:
    python3 embed_qtconsole.py
非阻塞的qtconsole,spyder用的就是这个。
添加了变量表格，以及变量传输。

"""
import sys
from qtconsole.rich_jupyter_widget import RichJupyterWidget
from qtconsole.manager import QtKernelManager
from PyQt5.Qt import QApplication
from PyQt5.QtWidgets import QMainWindow, QWidget, QHBoxLayout, QPushButton, QTableWidget, QTableWidgetItem

import ast
import base64
import pickle

# The ID of an installed kernel, e.g. 'bash' or 'ir'.
USE_KERNEL = 'python3'


class MyJupyterWidget(RichJupyterWidget):
    def __init__(self):
        super().__init__()
        self.callback = None

        self.on_data_got=None
        self.reset_message_id = 0

    def _init_command(self):
        cmd = """
import inspect
import base64
import pickle
import numpy as np
a=np.array([1,2,3,4,5]*1000)
def __save_vars__():
    __temp_dic = dict()
    globals_ = globals()
    for __k in globals_.keys():
        if not ( (__k.startswith('_')) or (__k in ['exit','quit','get_ipython','In','Out'])):
            if not inspect.ismodule(globals_[__k]):
                shape = None
                tmp = globals_[__k]
                if hasattr(tmp,'shape'):
                    shape= tmp.shape
                elif hasattr(tmp,'__len__'):
                    shape=len(tmp)
                else:
                    shape=1
                s = str(globals_[__k])
                if len(s)>100:
                    s=s[:100]
                __temp_dic[__k] = [str(shape),s]
    print('<<%s>>' % str(__temp_dic))
def __inspect__(__k:str):
    __data={'name':__k,'value':globals()[__k]}
    print('{{'+base64.b64encode(pickle.dumps(__data)).decode('ascii'))
def __add_var__(__name:str,__value_b64:str):
    globals()[__name]=pickle.loads(base64.b64decode(__value_b64))
        """
        self._execute(cmd, hidden=True)

    def _execute(self, source: str, hidden: bool):
        print('source code:',source)
        if not source.isidentifier():
            equal_symbol_index = source.find('=')
            cmp_sym_index = source.find('==')
            if equal_symbol_index != -1:
                if equal_symbol_index < cmp_sym_index or cmp_sym_index == -1:  # 确定是赋值符号，赋值符号必须在比较符号的左边。
                    source += """\n__save_vars__()"""

        msg_id = self.kernel_client.execute(source, hidden)
        print(source.find('%reset')!=-1,source.strip().startswith('reset'))
        print(source.find('%reset')!=-1 | source.strip().startswith('reset'))
        if (source.find('%reset')!=-1) or (source.strip().startswith('reset')):

            self.reset_message_id = msg_id

        print('sent msg id:',msg_id,self.reset_message_id)# 发送的命令的信息。

        self._request_info['execute'][msg_id] = self._ExecutionRequest(msg_id, 'user')
        self._hidden = hidden
        if not hidden:
            self.executing.emit(source)

    def _handle_execute_reply(self, msg):
        print('handle exe reply:',msg)
        print('msgid',msg['parent_header']['msg_id'])
        print(type(self.reset_message_id))
        if msg['parent_header']['msg_id']==self.reset_message_id:
            from threading import Timer
            t = Timer(0.2,self._init_command )
            t.start()
        super()._handle_execute_reply(msg)

    def _handle_execute_input(self, msg):
        super()._handle_execute_result(msg)

    def _handle_execute_result(self, msg):
        super()._handle_execute_result(msg)
        pass

    def set_var(self, name, value):
        source = """
__add_var__(\'{name}\',\'{value}\')
__save_vars__()
        """.format(name=name, value=base64.b64encode(pickle.dumps(value)).decode('ascii'))
        self._execute(source, hidden=True)

    def get_var(self, name: str):
        source = """
__inspect__(\'{name}\')
                """.format(name=name)
        self._execute(source, hidden=True)

    def _handle_stream(self, msg):  # 处理返回的信息。其中有
        content = msg['content']['text']  # 是数据中由print输出的内容。
        print('result content', content)
        print(msg)
        if content.find('<<') != -1:
            if content.find('{{') != -1:  # “{{”符号一旦出现，那么一定在“<<”之前。另外，一旦此符号出现，说明出现了base64字符。base64字符不包含{{。
                inspected_data = content.split('{{')[1]
                print('inspected data', inspected_data)
                msg['content']['text'] = content.split('{{')[0]
                import numpy as np
                data = pickle.loads(base64.b64decode(inspected_data))  # 直接可以获取对象了。
                print(data[666])
            else:
                data_dic = content.split('<<')[1][:-3]  # 化成数组的形状。
                print('data dic:', data_dic)
                msg['content']['text'] = content.split('<<')[0]
                if callable(self.callback):
                    self.callback(data_dic)
        elif content.find("{{") != -1:  # 一旦处于获取信息的状态。

            inspected_data = content.split('{{')[1]
            # print('inspected data', inspected_data)
            msg['content']['text'] = content.split('{{')[0]

            import numpy as np
            data = pickle.loads(base64.b64decode(inspected_data))  # 直接可以获取对象了。
            if callable(self.on_data_got):
                self.on_data_got(data)
            # 这里是获取信息的状态。

        super()._handle_stream(msg)
        pass


def make_jupyter_widget_with_kernel():
    """Start a kernel, connect to it, and create a RichJupyterWidget to use it
    """
    kernel_manager = QtKernelManager(kernel_name=USE_KERNEL)
    kernel_manager.start_kernel()

    kernel_client = kernel_manager.client()
    kernel_client.start_channels()

    jupyter_widget = MyJupyterWidget()
    jupyter_widget.kernel_manager = kernel_manager
    jupyter_widget.kernel_client = kernel_client
    return jupyter_widget


class IPythonWidget(QWidget):
    def __init__(self):
        super().__init__()

        self.jupyter_widget = make_jupyter_widget_with_kernel()
        self.jupyter_widget._init_command()
        self.table = QTableWidget()
        layout = QHBoxLayout()
        layout.addWidget(self.jupyter_widget)
        btn = QPushButton(text='hahahaha')
        layout.addWidget(btn)
        layout.addWidget(self.table)
        self.setLayout(layout)
        self.table.setRowCount(1)
        self.table.setColumnCount(3)
        self.table.setHorizontalHeaderLabels(['name', 'size', 'value'])
        btn.clicked.connect(self.conntestfunc)

        self.set_var('abv', 'ahhuhua' * 10)# 传入参数，应当可以传入任意类型
        self.get_var('abv')# 获取变量
        self.jupyter_widget.on_data_got = lambda data:self.handle_on_get_data(data)

        # lambda: self.jupyter_widget.set_var('a', [1] * 1024 * 1024))
        # print(self.jupyter_widget.kernel_client.get_connection_info())
        # 10MB大小的字符串,传输比较迅速。
        # 100MB大小的字符串也不至于卡死。

        self.jupyter_widget.callback = lambda msg: self.update_widget(msg=msg)

    def handle_on_get_data(self,data):
        print('got data,name : {name}, value : {value}'.format(name=data['name'],value=data['value']))

    def conntestfunc(self):
        '''
        zeromq的通信测试，目前无用。
        :return:
        '''
        print(self.jupyter_widget.kernel_client.get_connection_info())
        # msg1={'header': {'msg_id': 'ddd6b757-8aba820dda5b6a8f1cccde5c_11', 'msg_type': 'display_data',
        #                  'username': 'hzy', 'session': 'ddd6b757-8aba820dda5b6a8f1cccde5c', 'date': datetime(2020, 8, 19, 2, 8, 31, 181209, tzinfo=tzutc()), 'version': '5.3'},
        #       'msg_id': 'ddd6b757-8aba820dda5b6a8f1cccde5c_11', 'msg_type': 'stream',
        #       'parent_header': {'msg_id': '60431263-82d528ec6dd4012ff89e55a7_3', 'msg_type': 'execute_request',
        #                         'username': 'hzy', 'session': '60431263-82d528ec6dd4012ff89e55a7',
        #                         'date': datetime(2020, 8, 19, 2, 8, 30, 291533, tzinfo=tzutc()), 'version': '5.3'}
        #     , 'metadata': {}, 'content': {'name': 'stdout', 'text': "<<{'a': ['(5000,)', '[1 2 3 ... 3 4 5]']}>>\n"}, 'buffers': []}
        #
        # msg = {'metadata' : {'application/json' : {'expanded': True},},
        #        'data':'123',
        #     'transient': {'display_id': 'abcd'}
        # }
        ##############
        # content= dict(code='im', cursor_pos=len('im'))
        # msg = self.jupyter_widget.kernel_client.session.msg('complete_request',content)
        # print(msg)
        # self.jupyter_widget.kernel_client.shell_channel.send(msg)
        self.jupyter_widget.kernel_client.shell_channel.message_received.connect(
            lambda msg: print('shell-message', msg))
        self.jupyter_widget.kernel_client.iopub_channel.message_received.connect(lambda msg: print('iopub-msg', msg))
        self.jupyter_widget.kernel_client.stdin_channel.message_received.connect(lambda msg: print('stdin-msg', msg))
        ########
        content = {'data': {'a': '123'}, 'metadata': {'image/png': {'width': 640, 'height': 400}}
            , 'transient': {'display_id': 'abcd'}}
        msg = self.jupyter_widget.kernel_client.session.msg('display_data', content)
        self.jupyter_widget.kernel_client.iopub_channel.send(msg)

    def update_widget(self, msg):
        # print('message', msg)
        valid_json = ast.literal_eval(msg)  # json.loads(msg)
        self.table.setRowCount(len(list(valid_json.keys())))
        for i, k in enumerate(valid_json.keys()):
            name_item = QTableWidgetItem(k)
            self.table.setItem(i, 0, name_item)
            shape_item = QTableWidgetItem(valid_json[k][0])
            self.table.setItem(i, 1, shape_item)
            length = len(valid_json[k][1])
            if length >= 100:
                content = valid_json[k][1][:100]
            else:
                content = valid_json[k][1]
            value_item = QTableWidgetItem(content)
            self.table.setItem(i, 2, value_item)

    def set_var(self, name: str, value: object):
        '''
        向控制台中传递参数。
        :param name:
        :param value:
        :return:
        '''
        self.jupyter_widget.set_var(name, value)

    def get_var(self, name: str):
        '''
        获取控制台变量名
        :param name:
        :return:
        '''
        return self.jupyter_widget.get_var(name)

    def execute(self, cmd: str, hidden=True):
        '''
        :param cmd:str ,命令，合法的python命令
        :param hidden: 隐藏。为True的时候，执行时在界面上不会显示。为False的时候会显示。
        :return: None
        '''
        # self.jupyter_widget.append_stream("aaaaa")# 意思是
        # msg_id = self.jupyter_widget.kernel_client.execute(cmd, hidden)
        # a = self.jupyter_widget ._ExecutionRequest(msg_id, 'user')
        # print(a)
        self.jupyter_widget._execute(cmd, hidden=hidden)

    def shutdown_kernel(self):
        print('Shutting down kernel...')
        self.jupyter_widget.kernel_client.stop_channels()
        self.jupyter_widget.kernel_manager.shutdown_kernel()

class MainWindow(QMainWindow):
    def __init__(self):
        super().__init__()
        centralWidget = IPythonWidget()
        self.ipythonWidget = centralWidget
        self.setCentralWidget(centralWidget)

if __name__ == "__main__":
    app = QApplication(sys.argv)
    window = MainWindow()
    window.show()
    app.aboutToQuit.connect(window.ipythonWidget.shutdown_kernel)
    sys.exit(app.exec_())
